#include <iostream>
#include <opencv2/opencv.hpp>

#include "param.h"


void preprocess(const cv::Mat &src, cv::Mat &img) {
  cv::cvtColor(src, img, cv::COLOR_BGR2GRAY);
  cv::Mat brightened_frame;
  img.convertTo(brightened_frame, -1, param::alpha, param::beta);
  // 拉普拉斯锐化
  cv::Mat kernel =
      (cv::Mat_<float>(3, 3) << 0, -2, 0, -2, 9, -2, 0, -2, 0); // 锐化核
  cv::filter2D(brightened_frame, img, img.depth(), kernel);
}

// 创建 HOG 描述符实例
cv::HOGDescriptor hog(cv::Size(100, 100), // winSize - 必须与图像尺寸一致
                      cv::Size(16, 16),  // blockSize
                      cv::Size(4, 4),    // blockStride
                      cv::Size(8, 8),    // cellSize
                      9                  // nbins
);

// 存储特征向量和标签
std::vector<std::vector<float>> feature_vectors_list;
std::vector<int> labels_list;

void hogPreprocess(const std::string& path, int count, int id) {
  for (int i = 1; i <= count; ++i) {
    std::string img_path = path + std::to_string(i) + ".jpg";
    cv::Mat image = cv::imread(img_path);
    if (image.empty())
      continue;

    preprocess(image, image);
    // 计算 HOG 描述符
    std::vector<float> descriptors;
    // 最后的 0 表示不使用掩码
    hog.compute(image, descriptors, cv::Size(8, 8), cv::Size(0, 0));

    if (!descriptors.empty()) {
      feature_vectors_list.push_back(descriptors);
      labels_list.push_back(id);
    }
  }
}

int main() {
  const int count2 = 1000;
  const int count0 = 51;
  const int count6 = 800;
  const int count7 = 800;
  hogPreprocess("train/0/", count0, 0);
  hogPreprocess("train/2/", count2, 0);
  hogPreprocess("train/6/", count6, 0);
  hogPreprocess("train/7/", count7, 7);

  // 3. 转换为 SVM 所需的 cv::Mat 格式 (CV_32F)
  if (feature_vectors_list.empty()) {
    std::cerr << "Error: No features were successfully extracted." << std::endl;
    return -1;
  }

  // D (特征维度) 是第一个样本的特征向量长度
  int D = feature_vectors_list[0].size();
  int N = feature_vectors_list.size();

  // 创建最终的训练数据矩阵 (N x D)
  cv::Mat training_data_mat(N, D, CV_32F);
  cv::Mat labels_mat(N, 1, CV_32S); // 标签通常是 32 位整数

  // 填充矩阵
  for (int i = 0; i < N; ++i) {
    // 复制特征向量到训练数据矩阵的行
    std::memcpy(training_data_mat.ptr<float>(i), feature_vectors_list[i].data(),
                D * sizeof(float));

    // 填充标签
    labels_mat.at<int>(i, 0) = labels_list[i];
  }

  // 1. 设置训练参数
  cv::Ptr<cv::ml::SVM> svm = cv::ml::SVM::create();
  svm->setType(cv::ml::SVM::C_SVC);    // 分类模式
  svm->setKernel(cv::ml::SVM::LINEAR); // 线性核 (推荐用于HOG，计算快)
  // svm->setKernel(cv::ml::SVM::RBF);      // RBF核 (如果线性效果不好)

  // 3. 训练
  svm->train(training_data_mat, cv::ml::ROW_SAMPLE, labels_mat);

  // 4. 保存模型
  svm->save("7.xml");
}